package com.hoas.distributedlock.strategy.zookeeper;

import com.hoas.distributedlock.core.DisributedLock;
import com.hoas.distributedlock.core.ILockName;
import org.I0Itec.zkclient.ZkClient;

import java.util.HashSet;
import java.util.Set;
import java.util.concurrent.TimeUnit;

public class ZookeeperLock implements DisributedLock {

    private Set<String> locks = new HashSet<String>();
    private String prefix = "disributedlock";

    private ZkClient zkClient;

    public ZookeeperLock(ZkClient zkClient) {
        this.zkClient = zkClient;
    }

    public boolean lock(ILockName lockName, int times) {
        times = times < 1 ? 1: times;
        while (times > 0) {
            if (lock(lockName)){
                return true;
            }
            try {
                TimeUnit.MILLISECONDS.sleep(50);
            } catch (InterruptedException e) {
                e.printStackTrace();
            }

        }

        return false;
    }

    public boolean lock(ILockName lockName) {

        String key = lockName(lockName);
        // 控制重入
        if (locks.contains( key )) {
            return true;
        }

        if(zk_setnx(key)) {

            locks.add( key);
            return true;

        }

        return false;
    }

    public void release(ILockName lockName) {
        zk_del(lockName(lockName));
        locks.remove( lockName.value() );
    }

    public void release() {
        String[] keys = new String[locks.size()];
        locks.toArray(keys);
        zk_del(keys);
        zkClient.close();
        locks.clear();
    }

    public String getPrefix() {
        return prefix;
    }

    public void setPrefix(String prefix) {
        this.prefix = prefix;
    }

    protected String setData() {
        return "" + System.currentTimeMillis();
    }

    protected String lockName(ILockName lockName) {
        return prefix + "/" + lockName.value();
    }

    //尝试获得锁
    protected boolean zk_setnx(String key) {

        key = "/" + key.replaceAll(":", "/");

        try {
            zkClient.createEphemeral(key);
            return true;
        } catch (Exception e) {
            return false;
        }
    }

    protected void zk_del(String...keys) {

        for (int i = 0; i < keys.length; i++) {
            String key = "/" +  keys[i].replaceAll(":", "/");
            zkClient.delete( keys[i] );
        }

    }


}
